#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from kafka import KafkaProducer
import json
import os
import time


class KafkaMessageWriter:
    def __init__(self, bootstrap_servers='192.168.56.101:9092'):
        self.bootstrap_servers = bootstrap_servers
        self.producer = None
        self.topics = ['panda', 'bfd']

    def connect(self):
        """连接到 Kafka 集群"""
        try:
            self.producer = KafkaProducer(
                bootstrap_servers=self.bootstrap_servers,
                value_serializer=lambda v: json.dumps(v).encode('utf-8'),
                key_serializer=lambda v: json.dumps(v).encode('utf-8') if v else None,
                acks='all',
                retries=3
            )
            print(f"✅ 成功连接到 Kafka: {self.bootstrap_servers}")
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False

    def write_message(self):
        """交互式写入消息"""
        if not self.producer:
            print("❌ 请先连接到 Kafka")
            return

        # 选择 topic
        print("\n📋 可用 Topic:")
        for i, topic in enumerate(self.topics, 1):
            print(f"{i}. {topic}")

        try:
            choice = int(input("\n请选择 Topic 编号 (1-2): ")) - 1
            if choice < 0 or choice >= len(self.topics):
                print("❌ 无效的选择")
                return
            topic = self.topics[choice]
        except ValueError:
            print("❌ 请输入有效数字")
            return

        # 输入 key
        key_input = input("请输入 Key (直接回车表示空): ").strip()
        key = key_input if key_input else None

        # 输入 value
        while True:
            value = input("请输入 Value (不能为空): ").strip()
            if value:
                break
            print("❌ Value 不能为空，请重新输入")

        # 发送消息
        try:
            future = self.producer.send(
                topic=topic,
                key=key,
                value=value
            )
            # 等待发送确认
            record_metadata = future.get(timeout=10)
            print(f"✅ 消息发送成功!")
            print(f"   Topic: {record_metadata.topic}")
            print(f"   Partition: {record_metadata.partition}")
            print(f"   Offset: {record_metadata.offset}")
            if key:
                print(f"   Key: {key}")
            print(f"   Value: {value}")

        except Exception as e:
            print(f"❌ 消息发送失败: {e}")

    def upload_file(self):
        """上传文件内容"""
        if not self.producer:
            print("❌ 请先连接到 Kafka")
            return

        # 选择 topic
        print("\n📋 可用 Topic:")
        for i, topic in enumerate(self.topics, 1):
            print(f"{i}. {topic}")

        try:
            choice = int(input("\n请选择 Topic 编号 (1-2): ")) - 1
            if choice < 0 or choice >= len(self.topics):
                print("❌ 无效的选择")
                return
            topic = self.topics[choice]
        except ValueError:
            print("❌ 请输入有效数字")
            return

        # 输入文件路径
        file_path = input("请输入本地文件路径: ").strip()

        if not os.path.exists(file_path):
            print("❌ 文件不存在")
            return

        if not os.path.isfile(file_path):
            print("❌ 路径不是文件")
            return

        # 输入 key
        key_input = input("请输入 Key (直接回车表示空): ").strip()
        key = key_input if key_input else None

        # 读取并发送文件内容
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                line_count = 0
                success_count = 0

                print(f"\n📁 开始读取文件: {file_path}")

                for line_num, line in enumerate(file, 1):
                    line = line.strip()
                    if line:  # 跳过空行
                        line_count += 1
                        try:
                            # 可以为每行生成不同的 key
                            line_key = f"{key}_{line_num}" if key else None

                            future = self.producer.send(
                                topic=topic,
                                key=line_key,
                                value=line
                            )
                            future.get(timeout=5)  # 等待发送确认
                            success_count += 1

                            if line_count % 100 == 0:
                                print(f"📤 已发送 {line_count} 行...")

                        except Exception as e:
                            print(f"❌ 第 {line_num} 行发送失败: {e}")

                print(f"\n✅ 文件上传完成!")
                print(f"   总行数: {line_count}")
                print(f"   成功发送: {success_count}")
                print(f"   失败: {line_count - success_count}")

        except Exception as e:
            print(f"❌ 文件读取失败: {e}")

    def run(self):
        """主运行循环"""
        print("🚀 Kafka 消息写入工具")
        print("=" * 40)

        if not self.connect():
            return

        while True:
            print("\n请选择操作:")
            print("1. 📝 写入消息")
            print("2. 📁 上传文件")
            print("3. 🚪 退出")

            choice = input("\n请输入选择 (1-3): ").strip()

            if choice == '1':
                self.write_message()
            elif choice == '2':
                self.upload_file()
            elif choice == '3':
                print("👋 再见!")
                break
            else:
                print("❌ 无效选择，请重新输入")

        # 关闭连接
        if self.producer:
            self.producer.close()


if __name__ == "__main__":
    writer = KafkaMessageWriter()
    writer.run()
